import os
import random

from baseline_multilabel import load_chest_annotations




dataset_type = 'chest'  # 数据集类型
K_shot = 1  # 当前shot数量

filepath = os.path.join(os.getcwd(), './configs/baseline_multilabel.yaml')

trainval_txt = './data/MedFMC/chest/trainval.txt'  # 训练集图像列表文件路径

img_list = load_chest_annotations(trainval_txt)  # 加载训练集图像信息

N_way = 19  # 类别数量


def gen_support_set(img_list, N_way, K_shot):
    """生成用于few-shot学习任务的Chest数据集列表。

    Args:
        img_list (List[dict]): 包含每个图像的文件名和标签注释的输入数据集文件信息列表。
        N_way (int): 标签中的类别数量。
        K_shot (int): 每个类别的样本数量，如1-shot、5-shot和10-shot。

    Returns:
        List[List[Tuple[str, List[int]]]]: 包含支持集样本的列表。列表中的每个元素都是一个类别的样本列表，
        每个样本由文件名和标签组成。
    """

    support_set = []
    for cls_idx in range(N_way):
        imgs_per_cls = []
        for item in img_list:
            label = item['gt_label']
            if label[cls_idx] == 1:
                imgs_per_cls.append([item['filename'], label])
        if len(imgs_per_cls) <= K_shot:
            sample_set = imgs_per_cls
        else:
            random.shuffle(imgs_per_cls)
            sample_set = imgs_per_cls[:K_shot]

        support_set.append(sample_set)

    return support_set


support_set = gen_support_set(img_list, N_way, K_shot)  # 生成支持集样本

few_shot_lst = []
with open(f'chest_1-shot_train.txt', 'w') as f:
    for i, i_class in enumerate(support_set):
        for j_id in support_set[i]:
            j_pid, j_label = j_id
            k_line = j_pid
            few_shot_lst.append(j_id)
            for k, k_label in enumerate(j_label.tolist()):
                if k == 0:
                    sep_str = ' '
                elif k != 0:
                    sep_str = ','
                k_line += sep_str + str(k_label)
            f.write(k_line + '\n')